
#include "replication_pad2d_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"


namespace optiling {
    const uint32_t BLOCK_SIZE = 128;
    const uint32_t BUFFER_NUM = 2;
    static ge::graphStatus TilingFunc(gert::TilingContext* context)
    {
        // auto shape = context->GetInputTensor(0)->GetOriginShape();
        auto out_shape = context->GetOutputShape(0)->GetOriginShape();
        auto in_shape = context->GetInputTensor(0)->GetOriginShape();


        const int32_t dimNum = in_shape.GetDimNum();
        int32_t input_shape[100];
        for(int32_t i = 0; i < dimNum; i++){
            input_shape[i] = in_shape.GetDim(i);
        }

        int32_t output_shape[100];
        for(int32_t i = 0; i < dimNum; i++){
            output_shape[i] = out_shape.GetDim(i);
        }

        uint64_t ubSize;

        uint32_t typeLength = 0;

        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());

        ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);

    
        int32_t input_size = input_shape[dimNum-1]*input_shape[dimNum-2];
        int32_t output_size = output_shape[dimNum-1]*output_shape[dimNum-2];
        
        int32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
        int32_t outputNum = context->GetOutputShape(0)->GetStorageShape().GetShapeSize();

        int32_t mode;


        // if((input_size+output_size)*typeLength<ubSize) exit(-1);

        if(((input_size+output_size)*typeLength)+output_size*4<ubSize) 
        {
            mode = 0;
        }
        else
        {
            mode = 1;
        }
        

        ReplicationPad2dTilingData tiling;


        int32_t input_stride = input_size;
        int32_t output_stride = output_size;
        
        int32_t batch_number = inputNum/input_stride;

        printf("%d %d %d %d %d \n", input_stride, output_stride, inputNum, outputNum, batch_number);
        tiling.set_inputNum(inputNum);
        tiling.set_outputNum(outputNum);
        tiling.set_input_stride(input_stride);
        tiling.set_output_stride(output_stride);
        tiling.set_batch_number(batch_number);

        tiling.set_last_dim(input_shape[dimNum-1]);
        tiling.set_mode(mode);
        


        context->SetBlockDim(1);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;
        return ge::GRAPH_SUCCESS;
    }
}


namespace ge {
    static ge::graphStatus InferShape(gert::InferShapeContext* context)
    {
        const gert::Shape* x1_shape = context->GetInputShape(0);
        gert::Shape* y_shape = context->GetOutputShape(0);
        *y_shape = *x1_shape;
        return GRAPH_SUCCESS;
    }
}


namespace ops {
class ReplicationPad2d : public OpDef {
public:
    explicit ReplicationPad2d(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("paddings")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(ReplicationPad2d);
}
